✨ Description

Extracted and refined the dataset preparation script from #17.
Made it a command like train or convert.
Example call and config:

fast-llm prepare gpt_memmap --config foo.yaml

or

torchrun --standalone --nnodes 1 --nproc_per_node=1 --no_python \
    fast-llm prepare gpt_memmap --config foo.yaml

where foo.yaml contains:

output_path: /tmp/foo

loading_workers: 4
tokenize_workers: 4
saving_workers: 4

dataset:
  name_or_path: stas/openwebtext-10k

tokenizer:
  path: /tmp/SmolLM-135M/tokenizer.json

Run git clone https://huggingface.co/HuggingFaceTB/SmolLM-135M in tmp to get that tokenizer file.

This will produce:

/tmp/foo
├── downloaded_dataset
│   ├── cache-1e5559f36da9962e_00002_of_00004.arrow
│   ├── cache-1e5559f36da9962e_00003_of_00004.arrow
│   ├── cache-1e5559f36da9962e_00001_of_00004.arrow
│   ├── cache-1e5559f36da9962e_00000_of_00004.arrow
│   ├── data-00001-of-00004.arrow
│   ├── dataset_info.json
│   ├── data-00000-of-00004.arrow
│   ├── data-00002-of-00004.arrow
│   ├── data-00003-of-00004.arrow
│   ├── ok
│   └── state.json
├── shard_0_0.idx
├── shard_0_0.bin
└── fast_llm_dataset.json

with fast_llm_dataset.json reading:

{
    "datasets": [
        {
            "prefix": "shard_0_0",
            "num_documents": 10000,
            "num_tokens": 11569536,
            "weight": 1.0
        }
    ]
}

The downloaded_dataset can be deleted afterwards. It is not used by Fast-LLM.

🔍 Type of change

Select all that apply:

• 🐛 Bug fix (non-breaking change that addresses a specific issue)
• 🚀 New feature (non-breaking change that adds functionality)
• ⚠️ Breaking change (a change that could affect existing functionality)
• 📈 Performance improvement/optimization (improves speed, memory usage, or efficiency)
• 🛠️ Code refactor (non-functional changes that improve code readability, structure, etc.)
• 📦 Dependency bump (updates dependencies, including Dockerfile or package changes)
• 📝 Documentation change (updates documentation, including new content or typo fixes)
• 🔧 Infrastructure/Build change (affects build process, CI/CD, or dependencies)

📝 Changes

1. Fixed memory-mapped indexed dataset and added round-trip tests
2. Added prepare_dataset command
3. Simplified Dockerfile

✅ Checklist

Make sure the following tasks are completed before submitting the PR:

General:

• 📜 I have read and followed the contributing guidelines.
• 🎉 The functionality is complete, and I have tested the changes.
• 📝 I have updated the documentation if needed.
• ⚠️ The change does not introduce any new issues (e.g., runtime warnings, type checker errors, linting problems, unhandled edge cases).
• 🧩 I have commented my code, especially in hard-to-understand areas.

Dependencies and Configuration:

• 🐋 I have updated the Docker configuration or dependencies, if applicable.
• 🔄 I have ensured compatibility with the existing setup after dependency changes.

Testing:

• 🧪 I have added or updated tests to cover my changes.
• ✔️ New and existing tests pass locally with my changes.
• 🚦 I have tested these changes on GPUs and verified training stability.
• 🏋️ I have tested the changes on realistic training workloads, if applicable.

Performance Impact:

• 📊 I have run benchmarks where applicable to evaluate the performance impact.
• ✅ The benchmarks show no performance regression.
• 🚀 The benchmarks indicate a potential performance improvement.
• ⚠️ The benchmarks indicate a potential performance degradation.
• 📈 I have provided benchmark results and detailed any performance impact below, if applicable.

📊 Performance Impact Details

N/A

📝 Additional Notes

N/A